#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
修复huggingface_hub版本兼容性问题
解决cached_download导入错误
"""

import os
import sys
import subprocess
import importlib
import warnings

def check_huggingface_hub_version():
    """检查huggingface_hub版本"""
    try:
        from huggingface_hub import __version__
        version = __version__
        print(f"当前huggingface_hub版本: {version}")
        return version
    except ImportError:
        print("huggingface_hub未安装")
        return None

def check_cached_download_availability():
    """检查cached_download是否可用"""
    try:
        from huggingface_hub import cached_download
        print("✅ cached_download可用")
        return True
    except ImportError:
        print("❌ cached_download不可用")
        return False

def apply_compatibility_patch(verbose=True):
    """应用兼容性补丁
    
    Args:
        verbose: 是否显示详细日志
    """
    try:
        from huggingface_hub import hf_hub_download
        try:
            import huggingface_hub  # 只在需要动态添加属性时导入
        except ImportError as e:
            if verbose:
                print(f"❌ 无法应用兼容性补丁: {e}")
            return False
        # 定义兼容性函数
        def cached_download(*args, **kwargs):
            """兼容性函数，将cached_download重定向到hf_hub_download"""
            warnings.warn(
                "cached_download is deprecated, use hf_hub_download instead",
                DeprecationWarning,
                stacklevel=2
            )
            return hf_hub_download(*args, **kwargs)
        # 将cached_download添加到huggingface_hub模块
        if not hasattr(huggingface_hub, 'cached_download'):
            huggingface_hub.cached_download = cached_download
            if verbose:
                print("✅ 兼容性补丁已应用")
            return True
        else:
            if verbose:
                print("✅ 兼容性补丁已存在")
            return True
    except ImportError as e:
        if verbose:
            print(f"❌ 无法应用兼容性补丁: {e}")
        return False

def downgrade_huggingface_hub():
    """降级huggingface_hub到兼容版本"""
    print("正在降级huggingface_hub到兼容版本...")
    # 卸载当前版本
    try:
        subprocess.run([sys.executable, "-m", "pip", "uninstall", "huggingface_hub", "-y"], 
                      check=True, capture_output=True)
        print("✅ 已卸载当前版本")
    except subprocess.CalledProcessError as e:
        print(f"⚠️ 卸载时出现警告: {e}")
    # 安装兼容版本
    try:
        subprocess.run([sys.executable, "-m", "pip", "install", "huggingface_hub==0.24.5"], 
                      check=True, capture_output=True)
        print("✅ 已安装huggingface_hub==0.24.5")
        return True
    except subprocess.CalledProcessError as e:
        print(f"❌ 安装失败: {e}")
        return False

def run_fix():
    """运行完整的修复流程"""
    # 检查当前版本
    current_version = check_huggingface_hub_version()
    # 检查cached_download是否可用
    if check_cached_download_availability():
        print("✅ 无需修复，cached_download已可用")
        return True
    # 尝试应用兼容性补丁
    print("尝试应用兼容性补丁...")
    if apply_compatibility_patch(verbose=True):
        # 重新检查
        if check_cached_download_availability():
            print("✅ 兼容性补丁成功，cached_download现在可用")
            return True
    # 如果补丁失败，尝试降级
    print("兼容性补丁失败，尝试降级huggingface_hub...")
    if downgrade_huggingface_hub():
        # 重新检查
        if check_cached_download_availability():
            print("✅ 降级成功，cached_download现在可用")
            return True
    print("❌ 修复失败，请手动处理")
    return False

def main():
    """主函数 - 命令行入口点"""
    return run_fix()

if __name__ == "__main__":
    main() 